package net.tavion.angelboot.framework.logger.rr;

import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.server.reactive.ServerHttpResponse;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.Charset;

/**
 * @author tavion
 */
public class WebfluxServerResponseDecorator extends org.springframework.http.server.reactive.ServerHttpResponseDecorator {
    HttpLogNode logNotes;

    public WebfluxServerResponseDecorator(ServerHttpResponse delegate, HttpLogNode logNotes) {
        super(delegate);
        this.logNotes = logNotes;
    }

    /**
     * 基于netty,我这里需要显示的释放一次dataBuffer,但是slice出来的byte是不需要释放的,
     * 与下层共享一个字符串缓冲池,gateway过滤器使用的是nettyWrite类,会发生response数据多次才能返回完全。
     * 在 ServerHttpResponseDecorator 之后会释放掉另外一个refCount.
     */
    @Override
    public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {

        if (body instanceof Flux ) {
            Flux<? extends DataBuffer> fluxBody = (Flux<? extends DataBuffer>) body;
            Publisher<? extends DataBuffer> re = fluxBody.map(dataBuffer -> readAndReBuildBody(dataBuffer));
            return super.writeWith(re);
        }else if(body instanceof Mono){
            Mono<? extends DataBuffer> monoBody = (Mono<? extends DataBuffer>) body;
            Publisher<? extends DataBuffer> re = monoBody.map(dataBuffer -> readAndReBuildBody(dataBuffer));
            return super.writeWith(re);
        }

        return super.writeWith(body);

    }

    @Override
    public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
        return writeWith(Flux.from(body).flatMapSequential(p -> p));
    }

    private  DataBuffer readAndReBuildBody(DataBuffer dataBuffer){
        byte[] content = new byte[dataBuffer.readableByteCount()];
        // 数据读入数组
        dataBuffer.read(content);
        // 释放掉内存
        DataBufferUtils.release(dataBuffer);
        // 记录返回值
        String s = new String(content, Charset.forName("UTF-8"));

        logNotes.setBody(s);

        byte[] uppedContent = new String(content, Charset.forName("UTF-8")).getBytes();
        return this.bufferFactory().wrap(uppedContent);
    }
}
